import concurrent
import threading
import uuid
import traceback
import os
import tempfile
import time
import requests

from mcp.server import FastMCP
from mcp.types import TextContent

from .write_xiaohongshu import XiaohongshuPoster

mcp = FastMCP("xhs")
phone = os.getenv("phone", "")
path= os.getenv("json_path","/Users/bruce/")
slow_mode=os.getenv("slow_mode","False").lower() == "true"
SYNC_WAIT_SECONDS=int(os.getenv("SYNC_WAIT_SECONDS","120"))

# 简单内存任务表：仅用于单进程运行
JOBS: dict[str, dict] = {}

def _set_job(job_id: str, **kwargs):
    job = JOBS.get(job_id, {})
    job.update(kwargs)
    JOBS[job_id] = job

def _start_job(target, *args, **kwargs) -> str:
    job_id = str(uuid.uuid4())
    _set_job(job_id, status="pending", result=None, error=None)
    def runner():
        _set_job(job_id, status="running")
        try:
            res = target(*args, **kwargs)
            _set_job(job_id, status="done", result=res, error=None)
        except Exception:
            _set_job(job_id, status="error", error=traceback.format_exc(), result=None)
    th = threading.Thread(target=runner, daemon=True)
    th.start()
    return job_id

def _wait_for_job(job_id: str, timeout_sec: int) -> tuple[bool, str]:
    """轮询等待后台任务完成。在超时时间内返回 (finished, message)。"""
    start = time.time()
    while time.time() - start < timeout_sec:
        job = JOBS.get(job_id)
        if job:
            status = job.get("status")
            if status == "done":
                return True, str(job.get("result"))
            if status == "error":
                return True, f"error: {job.get('error')}"
        time.sleep(1)
    return False, f"timeout waiting job {job_id}"

@mcp.tool()
def ping(message: str = "ping") -> list[TextContent]:
    """Quick health check tool: returns immediately with a pong message."""
    return [TextContent(type="text", text=f"pong: {message}")]
def login():
    poster = XiaohongshuPoster(path)
    poster.login(phone)
    time.sleep(1)
    poster.close()

def download_image(url):
    local_filename = url.split('/')[-1]
    temp_dir = tempfile.gettempdir()

    local_path = os.path.join(temp_dir, local_filename)  # 假设缓存地址为/tmp
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(local_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    return local_path

def download_images_parallel(urls):
    """
    并行下载图片到本地缓存地址
    """
    with concurrent.futures.ThreadPoolExecutor() as executor:
        results = list(executor.map(download_image, urls))
    return results

@mcp.tool()
def create_note(title: str, content: str, images: list, account: str | None = None) -> list[TextContent]:
    """Create a note asynchronously and return a job_id to avoid RPC timeouts."""
    def job_func():
        poster = XiaohongshuPoster(path, account=account)
        try:
            if len(images)>0 and isinstance(images[0], str) and images[0].startswith("http"):
                local_images = download_images_parallel(images)
            else:
                local_images = images
            code, info = poster.login_to_publish(title, content, local_images, slow_mode)
            poster.close()
            return info
        finally:
            try:
                poster.close()
            except Exception:
                pass
    job_id = _start_job(job_func)
    return [TextContent(type="text", text=f"job_id={job_id}")]


@mcp.tool()
def create_note_async(title: str, content: str, images: list, account: str | None = None) -> list[TextContent]:
    """Start creating a note asynchronously. Returns a job_id for polling.

    Args:
        title: the title of the note (post), which should not exceed 20 words
        content: the description of the note (post).
        images: the list of image paths or URLs to be included in the note (post)
    """
    def job_func():
        poster = XiaohongshuPoster(path=r"D://xhs", account=account)
        try:
            if len(images)>0 and isinstance(images[0], str) and images[0].startswith("http"):
                local_images = download_images_parallel(images)
            else:
                local_images = images
            code, info = poster.login_to_publish(title, content, local_images, slow_mode)
            poster.close()
            return info
        finally:
            try:
                poster.close()
            except Exception:
                pass

    job_id = _start_job(job_func)
    return [TextContent(type="text", text=f"job_id={job_id}")]


@mcp.tool()
def create_note_wait(title: str, content: str, images: list, wait_seconds: int | None = None, account: str | None = None) -> list[TextContent]:
    """Create note but wait for completion up to wait_seconds (default from SYNC_WAIT_SECONDS)."""
    # 复用异步任务，随后进行同步等待
    resp = create_note_async(title, content, images, account=account)
    # resp[0].text like "job_id=..."
    text = resp[0].text
    job_id = text.split("job_id=")[-1]
    timeout = wait_seconds if isinstance(wait_seconds, int) and wait_seconds > 0 else SYNC_WAIT_SECONDS
    finished, msg = _wait_for_job(job_id, timeout)
    if finished:
        return [TextContent(type="text", text=msg)]
    return [TextContent(type="text", text=f"waiting: {msg}; job_id={job_id}")]


@mcp.tool()
def create_video_note(title: str, content: str, videos: list, account: str | None = None) -> list[TextContent]:
    """Create a video note asynchronously and return a job_id to avoid RPC timeouts."""
    def job_func():
        poster = XiaohongshuPoster(path, account=account)
        try:
            if len(videos)>0 and isinstance(videos[0], str) and videos[0].startswith("http"):
                local_videos = download_images_parallel(videos)
            else:
                local_videos = videos
            code, info = poster.login_to_publish_video(title, content, local_videos, slow_mode)
            poster.close()
            return info
        finally:
            try:
                poster.close()
            except Exception:
                pass
    job_id = _start_job(job_func)
    return [TextContent(type="text", text=f"job_id={job_id}")]


@mcp.tool()
def create_video_note_async(title: str, content: str, videos: list, account: str | None = None) -> list[TextContent]:
    """Start creating a video note asynchronously. Returns a job_id for polling."""
    def job_func():
        poster = XiaohongshuPoster(path, account=account)
        try:
            if len(videos)>0 and isinstance(videos[0], str) and videos[0].startswith("http"):
                local_videos = download_images_parallel(videos)
            else:
                local_videos = videos
            code, info = poster.login_to_publish_video(title, content, local_videos, slow_mode)
            poster.close()
            return info
        finally:
            try:
                poster.close()
            except Exception:
                pass

    job_id = _start_job(job_func)
    return [TextContent(type="text", text=f"job_id={job_id}")]


@mcp.tool()
def create_video_note_wait(title: str, content: str, videos: list, wait_seconds: int | None = None, account: str | None = None) -> list[TextContent]:
    """Create video note but wait for completion up to wait_seconds (default from SYNC_WAIT_SECONDS)."""
    resp = create_video_note_async(title, content, videos, account=account)
    text = resp[0].text
    job_id = text.split("job_id=")[-1]
    timeout = wait_seconds if isinstance(wait_seconds, int) and wait_seconds > 0 else SYNC_WAIT_SECONDS
    finished, msg = _wait_for_job(job_id, timeout)
    if finished:
        return [TextContent(type="text", text=msg)]
    return [TextContent(type="text", text=f"waiting: {msg}; job_id={job_id}")]


@mcp.tool()
def get_job_status(job_id: str) -> list[TextContent]:
    """Check async job status. Returns: pending|running|done|error and message."""
    job = JOBS.get(job_id)
    if not job:
        return [TextContent(type="text", text=f"not_found: {job_id}")]
    status = job.get("status")
    if status == "done":
        return [TextContent(type="text", text=f"done: {job.get('result')}")]
    if status == "error":
        return [TextContent(type="text", text=f"error: {job.get('error')}")]
    return [TextContent(type="text", text=status)]

def main():
    mcp.run()

if __name__ == "__main__":
    main()